#!/usr/bin/env python
# encoding: utf-8
"""
@author: youfeng
@email: youfeng243@163.com
@license: Apache Licence
@file: gsxt_master_scheduler.py
@time: 2018/1/21 21:02
"""
# 日志模块
import time

from common.mongo import MongDb
from common.table_manage import TableManage
from config.beanstalk_conf import BEANSTALK_CONF
from config.mongo_conf import MONGO_DB_SOURCE
from config.schedule_conf import SCHEDULE_LIST
from ext.scheduler_handler import SchedulerHandler
from logger import Logger
from model.crawl_model import STORE_TYPE_LIST, STORE_TYPE_DETAIL, STORE_TYPE_ANNUAL

global_log = Logger('gsxt_master_scheduler.log').get_logger()


class ScanIter(object):
    # 最长失效时间
    MAX_INVALID_TIME = 300

    def __init__(self, db, table_name, select, log):
        self.db = db
        self.table_name = table_name
        self.select = select
        self.init_times = 0
        self.last_time = int(time.time())
        self.__iter = self.init_iter()
        self.log = log

    def init_iter(self):
        self.init_times += 1
        self.last_time = int(time.time())
        return self.db.traverse_batch(self.table_name, self.select)

    def get_item(self):
        now_time = time.time()
        while True:
            try:
                item = self.__iter.next()
            except StopIteration:
                item = None
            break

        if item is not None:
            self.last_time = now_time
        else:
            if now_time - self.last_time >= self.MAX_INVALID_TIME:
                self.__iter = self.init_iter()
                self.log.info("重新初始化迭代器: table_name = {} times = {}".format(
                    self.table_name, self.init_times))

        return item


# 调度类
class MasterScheduler(object):
    MAX_BEANSTALK_NUM = 20

    def __init__(self, beanstalk_conf, schedule_list, log):
        self.log = log
        # # 初始化消息队列
        # self.beanstalk_thread = BeanstalkHandler(beanstalk_conf=beanstalk_conf, log=log)
        # self.beanstalk_thread.start()

        # # 额外开启一个消息队列
        # self.beanstalk = PyBeanstalk(beanstalk_conf['host'], beanstalk_conf['port'])

        # 初始化mongodb
        self.source_db = MongDb(MONGO_DB_SOURCE['host'], MONGO_DB_SOURCE['port'], MONGO_DB_SOURCE['db'],
                                MONGO_DB_SOURCE['username'],
                                MONGO_DB_SOURCE['password'], log=log)

        # 调度列表
        self.schedule_list = schedule_list

        # 获取所有的tube
        self.scheduler_handler = SchedulerHandler(beanstalk_conf, self.get_tube_set(schedule_list), log=log)

        # 确认表索引
        self.create_index()

        # 初始化数据库迭代器
        self.db_iter_list = self.init_iter_list()

        self.log.info("初始化工商调度完成...")

    # 获得消息队列集合信息
    @staticmethod
    def get_tube_set(schedule_list):
        tube_set = set()
        if not isinstance(schedule_list, list):
            return tube_set

        for item in schedule_list:
            if not isinstance(item, dict):
                continue

            tube = item.get('tube')
            if tube is None:
                continue

            tube_set.add(tube)

        return tube_set

    # 确保索引存在
    def create_index(self):
        for item in self.schedule_list:
            index_list = item.get("index")
            if not isinstance(index_list, list):
                continue

            province = item.get('province')
            if province is None:
                province = 'gsxt'

            table_name = item.get('table_name')
            if table_name is None:
                table_name = TableManage.get_seed_table(province)

            if not isinstance(table_name, basestring):
                continue

            for index in index_list:
                self.log.info("开始创建索引: table_name = {} index = {}".format(table_name, index))
                if not isinstance(index, basestring):
                    self.log.error("当前索引错误: table_name = {} index = {}".format(table_name, index))
                    continue

                self.source_db.create_index(table_name, [(index, MongDb.ASCENDING)])
                self.log.info("创建索引完成: table_name = {} index = {}".format(table_name, index))

    # 初始化数据扫描迭代器
    def init_iter_list(self):
        iter_list = []
        for item in self.schedule_list:
            table_name = item.get('table_name')
            if table_name is None:
                iter_list.append(None)
                continue
            select = item.get('select')
            if isinstance(select, dict):
                select_list = [select]
            elif isinstance(select, list):
                select_list = select
            else:
                select_list = [{}]

            scan_list = []
            for select_item in select_list:
                scan_list.append(ScanIter(self.source_db, table_name, select_item, self.log))

            iter_list.append(scan_list)

        return iter_list

    # 获得当前需要被调度的企业
    def get_current_company(self, scan_list):

        item = self.get_current_item(scan_list)
        if item is not None:
            company = item.get('_id')
            return company

        return None

    # 获取当前ITEM
    @staticmethod
    def get_current_item(scan_list):
        for db_iter in scan_list:
            item = db_iter.get_item()
            if not isinstance(item, dict):
                continue

            return item

        return None

    # # 判断是否阻塞
    # def is_beanstalk_block(self, tube):
    #     try:
    #         count = self.beanstalk.get_tube_count(tube)
    #     except CommandFailed, (_, status, results):
    #         if status == 'NOT_FOUND':
    #             self.log.error('获取消息队列不存在: tube = {}'.format(tube))
    #             count = 0
    #         else:
    #             self.log.error('获取消息队列异常: tube = {} status = {}'.format(tube, status))
    #             count = self.MAX_BEANSTALK_NUM
    #     except Exception as e:
    #         count = self.MAX_BEANSTALK_NUM
    #         self.log.error("获取消息队列大小失败: tube = {}".format(tube))
    #         self.log.exception(e)
    #
    #     if count >= self.MAX_BEANSTALK_NUM:
    #         return True
    #
    #     return False

    def select_data_by_field(self, province, seed_table, field, scan_list):
        if field == STORE_TYPE_LIST or field == STORE_TYPE_DETAIL:
            company = self.get_current_company(scan_list)
            if company is None:
                return None, None

            send_item = {
                'province': province,
                'company': company,
                'seed_table': seed_table,
                'field': field
            }
            return send_item, company

        # 如果是年报
        if field == STORE_TYPE_ANNUAL:
            # {
            #     # 年报关键ID
            #     '_id': anCheId,
            #     # 解析出来的企业名称
            #     'company': parse_comany,
            #     'year': 2014,
            #     # 企业类型，不是必要的但需要保存
            #     'ent_type': 18
            #     # 年报类型，不是必要的但需要保存
            #     'annual_type': 'e',
            #     # 通过 ent_type 与 annual_type 计算出来的
            #     'base_href': xxx,
            #     # 通过 ent_type 与 annual_type 计算出来的
            #     'index_href': xxxx,
            #     '_in_time': xxx
            #     'crawl_status': 0, 1, 2
            # 'crawl_status_time': '2018-01-23'
            # }

            item = self.get_current_item(scan_list)
            if item is None:
                return None, None

            company = item.get('company')
            if company is None:
                return None, None

            year = item.get('year')
            if year is None:
                return None, None

            base_href = item.get('base_href')
            if base_href is None:
                return None, None

            index_href = item.get('index_href')
            if index_href is None:
                return None, None

            an_che_id = item.get('_id')
            if an_che_id is None:
                return None, None

            ent_type = item.get('ent_type')
            if ent_type is None:
                return None, None

            annual_type = item.get('annual_type')
            if annual_type is None:
                return None, None

            send_item = {
                'province': province,
                'company': company,
                'year': year,
                'base_href': base_href,
                'index_href': index_href,
                'an_che_id': an_che_id,
                'ent_type': ent_type,
                'annual_type': annual_type,
            }
            return send_item, company

        self.log.info("暂不支持解析当前属性: province = {} field = {} seed_table = {}".format(
            province, field, seed_table))
        return None, None

    # 开始调度抓取数据
    def start(self):
        count = 0
        self.log.info("启动调度...")
        while True:
            is_send_success = False
            for index, item in enumerate(self.schedule_list):
                tube = item.get('tube')
                if tube is None:
                    continue

                province = item.get('province')
                if province is None:
                    province = 'gsxt'

                seed_table = item.get('table_name')
                if seed_table is None:
                    seed_table = TableManage.get_seed_table(province)

                # 没指定类型则不调度
                field = item.get('field')
                if field is None:
                    self.log.error("没有指定field, index = {} province = {}".format(index, province))
                    continue

                # 判断内存队列是否可以发送数据
                if self.scheduler_handler.is_mem_queue_block(tube):
                    continue

                send_item, company = self.select_data_by_field(province, seed_table, field, self.db_iter_list[index])
                if send_item is None:
                    continue

                # # 获得当前调度表的数据
                # company = self.get_current_company(self.db_iter_list[index])
                # if company is None:
                #     continue

                count += 1

                # 不阻塞发送数据，允许数据丢失，下次重试
                is_send = self.scheduler_handler.push_msg(tube, send_item)

                # is_send = self.scheduler_handler.push_msg(tube, {
                #     'province': province,
                #     'company': company,
                #     'seed_table': seed_table,
                #     'field': field
                # })

                if not is_send:
                    self.log.warn("当前消息队列已满,丢弃数据: province:{} tube:{} company:{}".format(
                        province, tube, company))
                    continue

                is_send_success = True
                self.log.info("当前调度企业: count:{} province:{} field:{} company:{}".format(
                    count, province, field, company))

            if is_send_success:
                continue

            # 如果所有表都没有发送数据 则休眠
            time.sleep(1)


def main():
    MasterScheduler(BEANSTALK_CONF, SCHEDULE_LIST, global_log).start()


if __name__ == '__main__':
    main()
